package io;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.text.DecimalFormat;
import java.util.BitSet;
import pac.SearchDatastructure;
import statistics.SignificanceHelper;

public class OutputWriter {

	private SearchDatastructure searcher;
	private SignificanceHelper significanceHelper;
	private BufferedWriter bwriter;
	private DecimalFormat doubleF;
	private DecimalFormat twiceDoubleF;
	private DecimalFormat scientificF;
	
	public OutputWriter(SearchDatastructure I, SignificanceHelper S) throws IOException {
		searcher = I;
		significanceHelper = S;
		FileWriter fstream = new FileWriter( searcher.userParams.getStage2Output(), true );
		bwriter = new BufferedWriter (fstream);
		doubleF	= new DecimalFormat("#.##"); 
		twiceDoubleF	= new DecimalFormat("#.####"); 
		scientificF	= new DecimalFormat("0.##E00");
	}


	private String getCaseCarrierIDs ( int[] variableArray) {

		BitSet var1Cases = new BitSet();
		BitSet var2Cases, var3Cases, var4Cases;

		switch(variableArray.length) {
		case 1: 
			var1Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[0]);
			break;

		case 2: 
			var1Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[0]);
			var2Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[1]);
			var1Cases.and(var2Cases); //co-carriers
			break;

		case 3:
			var1Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[0]);
			var2Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[1]);
			var3Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[2]);
			var1Cases.and(var2Cases); //co-carriers
			var1Cases.and(var3Cases);
			break;

		default:
			var1Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[0]);
			var2Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[1]);
			var3Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[2]);
			var4Cases = searcher.snpIndices.getCaseCarrierIndex(variableArray[3]);
			var1Cases.and(var2Cases); //co-carriers
			var1Cases.and(var3Cases);
			var1Cases.and(var4Cases);
			break;
		}

		StringBuffer caseCarrierIDs = new StringBuffer();
		for (int i = var1Cases.nextSetBit(0); i >= 0; i = var1Cases.nextSetBit(i+1)) 
			caseCarrierIDs.append(searcher.caseInfo[i].getIND_ID()+",");

		return caseCarrierIDs.toString();	
	}

	private String getControlCarrierIDs ( int[] variableArray) {

		BitSet var1Controls = new BitSet();
		BitSet var2Controls, var3Controls, var4Controls;

		switch(variableArray.length) {
		case 1: 
			var1Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[0]);
			break;

		case 2: 
			var1Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[0]);
			var2Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[1]);
			var1Controls.and(var2Controls); //co-carriers
			break;

		case 3:
			var1Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[0]);
			var2Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[1]);
			var3Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[2]);
			var1Controls.and(var2Controls); //co-carriers
			var1Controls.and(var3Controls);
			break;

		default:
			var1Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[0]);
			var2Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[1]);
			var3Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[2]);
			var4Controls = searcher.snpIndices.getControlCarrierIndex(variableArray[3]);
			var1Controls.and(var2Controls); //co-carriers
			var1Controls.and(var3Controls);
			var1Controls.and(var4Controls);
			break;
		}

		StringBuffer controlCarrierIDs = new StringBuffer();
		for (int i = var1Controls.nextSetBit(0); i >= 0; i = var1Controls.nextSetBit(i+1)) 
			controlCarrierIDs.append(searcher.controlInfo[i].getIND_ID()+",");

		return controlCarrierIDs.toString();	
	}


	private String getSNPDetailsfromVariables(int[] variableArray) {
		StringBuffer interactionString = new StringBuffer();
		
		for(int varNo=0; varNo<2; varNo++) {				
			int snpNo = (int) Math.floor(variableArray[varNo]/2);
			char mode = ( variableArray[varNo] % 2 == 0 ) ? 'd' : 'r';						
			interactionString.append( (snpNo+1) + "," + mode + "," + searcher.snpList[snpNo].toString() + "\t" );	
		}
		
		return interactionString.toString();
	}
	
	
	private double[][][] computeExpectations(int[][][] observations)
	{ 
		int total = observations[1][1][1] + observations[1][1][0] + observations[1][0][1] + observations[1][0][0];
		double P = (double) ( observations[1][1][1] + observations[1][1][0] ) / total;		
		double Q = (double) ( observations[1][1][1] + observations[1][0][1] ) / total;
		
		double[][][] expectations = new double[2][2][2];
		expectations[1][1][1] = P * Q * total;
		expectations[1][1][0] = P * (1-Q) * total;
		expectations[1][0][1] = (1-P) * Q * total;
		expectations[1][0][0] = (1-P) * (1-Q) * total;
		
		return expectations;
	}
	
	/*
	 * returns an array of 2 z-scores - case disequilibrium and control disequilibrium
	 */
	private double[] getCaseAndControlDisequilibria ( 
			double[][][] expectedCaseCarriers, int[][][] noOfCaseCarriers, 
			double[][][] expectedControlCarriers, int[][][] noOfControlCarriers ) 
	{
		//default, can't compute because of div-by-zero
		double Dprime_cases = -999;
		double Dprime_controls = -999;
		
		if ( expectedCaseCarriers[1][1][1] > 0 && expectedCaseCarriers[1][0][0] > 0
				&& expectedControlCarriers[1][1][1] > 0 && expectedControlCarriers[1][0][0] > 0 ) 
		{
			double P = ( expectedCaseCarriers[1][1][1] + expectedCaseCarriers[1][1][0] ) / searcher.userParams.getNumber_of_cases();		
			double Q = ( expectedCaseCarriers[1][1][1] + expectedCaseCarriers[1][0][1] ) / searcher.userParams.getNumber_of_cases();
			//double sigma_cases  = Math.sqrt( P*(1-P)*Q*(1-Q) / searcher.userParams.getNumber_of_cases() );
			double Dcases = (noOfCaseCarriers[1][1][1] -  expectedCaseCarriers[1][1][1]) / searcher.userParams.getNumber_of_cases() ;
			double Dcases_max = (Dcases <= 0) ? Math.min( P*Q, (1-P)*(1-Q) ) : Math.min( P*(1-Q), (1-P)*Q );
			//z_cases = D_cases / sigma_cases;
			Dprime_cases = Dcases / Dcases_max;
			
			double p = ( expectedControlCarriers[1][1][1] + expectedControlCarriers[1][1][0] ) / searcher.userParams.getNumber_of_controls();
			double q = ( expectedControlCarriers[1][1][1] + expectedControlCarriers[1][0][1] ) / searcher.userParams.getNumber_of_controls();
			//double sigma_controls = Math.sqrt( p*(1-p)*q*(1-q) / searcher.userParams.getNumber_of_controls() );
			double Dcontrols = (noOfControlCarriers[1][1][1] -  expectedControlCarriers[1][1][1]) / searcher.userParams.getNumber_of_controls();
			double Dcontrols_max = (Dcontrols <= 0) ? Math.min( p*q, (1-p)*(1-q) ) : Math.min( p*(1-q), (1-p)*q );
			//z_controls = D_controls / sigma_controls;
			Dprime_controls = Dcontrols / Dcontrols_max;
		}
		
		return new double[] {Dprime_cases, Dprime_controls};		
	}
	
	
	private synchronized void WriteOutputFull ( int[] variableArray, 
			double[][][] expectedCaseCarriers, int[][][] noOfCaseCarriers, double pvalueCaseOnly, 
			double[][][] expectedControlCarriers, int[][][] noOfControlCarriers, double pvalueCaseControl 
			) throws IOException {

		int order = variableArray.length;
		String interactionString = getSNPDetailsfromVariables(variableArray);
		double[] disequilibria = getCaseAndControlDisequilibria(expectedCaseCarriers, noOfCaseCarriers, expectedControlCarriers, noOfControlCarriers);
		double[] pvalueCaseOnly_both = significanceHelper.asymptoticAndExact_CaseOnly ( order, noOfCaseCarriers );
		double[] pvalueCaseControl_both = significanceHelper.asymptoticAndExact_CaseVsControl ( order, noOfCaseCarriers, noOfControlCarriers );
		
		String caseIDs = getCaseCarrierIDs(variableArray);
		String controlIDs = getControlCarrierIDs(variableArray);
		
		bwriter.write ( interactionString.toString() 
				+ doubleF.format(expectedCaseCarriers[1][0][0]) + "/" + doubleF.format(expectedCaseCarriers[1][0][1]) + "/" + doubleF.format(expectedCaseCarriers[1][1][0]) + "/" + doubleF.format(expectedCaseCarriers[1][1][1]) + "\t" 
				+ noOfCaseCarriers[1][0][0]+"/"+noOfCaseCarriers[1][0][1]+"/"+noOfCaseCarriers[1][1][0]+"/"+noOfCaseCarriers[1][1][1] +"\t" 					
				+ scientificF.format(pvalueCaseOnly_both[0]) + "\t" 
				+ scientificF.format(pvalueCaseOnly_both[1]) + "\t" 
				+ doubleF.format(expectedControlCarriers[1][0][0]) + "/" + doubleF.format(expectedControlCarriers[1][0][1]) + "/" + doubleF.format(expectedControlCarriers[1][1][0]) + "/" + doubleF.format(expectedControlCarriers[1][1][1]) + "\t" 
				+ noOfControlCarriers[1][0][0]+"/"+noOfControlCarriers[1][0][1]+"/"+ noOfControlCarriers[1][1][0]+"/"+noOfControlCarriers[1][1][1]+ "\t"
				+ scientificF.format(pvalueCaseControl_both[0]) + "\t"
				+ scientificF.format(pvalueCaseControl_both[1]) + "\t"
				+ twiceDoubleF.format(disequilibria[0]) + "\t"
				+ twiceDoubleF.format(disequilibria[1]) + "\t"
				+ caseIDs + "\t"
				+ controlIDs + "\n" ); 

		bwriter.flush();
	}

	private synchronized void WriteOutputNormal ( int[] variableArray, 
			double[][][] expectedCaseCarriers, int[][][] noOfCaseCarriers, double pvalueCaseOnly, 
			double[][][] expectedControlCarriers, int[][][] noOfControlCarriers, double pvalueCaseControl 
			) throws IOException { 
		
		int order = variableArray.length;
		String interactionString = getSNPDetailsfromVariables(variableArray);
		double[] disequilibria = getCaseAndControlDisequilibria(expectedCaseCarriers, noOfCaseCarriers, expectedControlCarriers, noOfControlCarriers);
		double[] pvalueCaseOnly_both = significanceHelper.asymptoticAndExact_CaseOnly ( order, noOfCaseCarriers );
		double[] pvalueCaseControl_both = significanceHelper.asymptoticAndExact_CaseVsControl ( order, noOfCaseCarriers, noOfControlCarriers );

		bwriter.write ( interactionString.toString() 
				+ doubleF.format(expectedCaseCarriers[1][0][0]) + "/" + doubleF.format(expectedCaseCarriers[1][0][1]) + "/" + doubleF.format(expectedCaseCarriers[1][1][0]) + "/" + doubleF.format(expectedCaseCarriers[1][1][1]) + "\t" 
				+ noOfCaseCarriers[1][0][0]+"/"+noOfCaseCarriers[1][0][1]+"/"+noOfCaseCarriers[1][1][0]+"/"+noOfCaseCarriers[1][1][1] +"\t" 					
				+ scientificF.format(pvalueCaseOnly_both[0]) + "\t" 
				+ scientificF.format(pvalueCaseOnly_both[1]) + "\t" 
				+ doubleF.format(expectedControlCarriers[1][0][0]) + "/" + doubleF.format(expectedControlCarriers[1][0][1]) + "/" + doubleF.format(expectedControlCarriers[1][1][0]) + "/" + doubleF.format(expectedControlCarriers[1][1][1]) + "\t" 
				+ noOfControlCarriers[1][0][0]+"/"+noOfControlCarriers[1][0][1]+"/"+ noOfControlCarriers[1][1][0]+"/"+noOfControlCarriers[1][1][1]+ "\t"
				+ scientificF.format(pvalueCaseControl_both[0]) + "\t"
				+ scientificF.format(pvalueCaseControl_both[1]) + "\t"
				+ twiceDoubleF.format(disequilibria[0]) + "\t"
				+ twiceDoubleF.format(disequilibria[1]) + "\n" ); 

		bwriter.flush();
	}
	
	
	private synchronized void WriteOutputBrief ( int[] variableArray, 
			double[][][] expectedCaseCarriers, int[][][] noOfCaseCarriers, double pvalueCaseOnly, 
			double[][][] expectedControlCarriers, int[][][] noOfControlCarriers, double pvalueCaseControl 
			) throws IOException { 
	
		String interactionString = getSNPDetailsfromVariables(variableArray);
		double[] disequilibria = getCaseAndControlDisequilibria(expectedCaseCarriers, noOfCaseCarriers, expectedControlCarriers, noOfControlCarriers);

		bwriter.write( interactionString.toString() +
				doubleF.format(expectedCaseCarriers[1][0][0]) + "/" + doubleF.format(expectedCaseCarriers[1][0][1]) + "/" + doubleF.format(expectedCaseCarriers[1][1][0]) + "/" + doubleF.format(expectedCaseCarriers[1][1][1]) + "\t" 
				+ noOfCaseCarriers[1][0][0]+"/"+noOfCaseCarriers[1][0][1]+"/"+noOfCaseCarriers[1][1][0]+"/"+noOfCaseCarriers[1][1][1] +"\t" 					
				+ scientificF.format(pvalueCaseOnly) + "\t" 
				+ doubleF.format(expectedControlCarriers[1][0][0]) + "/" + doubleF.format(expectedControlCarriers[1][0][1]) + "/" + doubleF.format(expectedControlCarriers[1][1][0]) + "/" + doubleF.format(expectedControlCarriers[1][1][1]) + "\t" 
				+ noOfControlCarriers[1][0][0]+"/"+noOfControlCarriers[1][0][1]+"/"+ noOfControlCarriers[1][1][0]+"/"+noOfControlCarriers[1][1][1]+ "\t"
				+ scientificF.format(pvalueCaseControl) + "\t" 
				+ twiceDoubleF.format(disequilibria[1]) + "\t"
				+ twiceDoubleF.format(disequilibria[0]) + "\n" );  
	
		bwriter.flush();
	}
	
	
	public void WriteOutput( int[] variableArray,
												int[][][] noOfCaseCarriers, double pvalueCaseOnly, 
												int[][][] noOfControlCarriers, double pvalueCaseControl ) 
	{
		
		double[][][] expectedCaseCarriers = computeExpectations(noOfCaseCarriers);
		double[][][] expectedControlCarriers = computeExpectations(noOfControlCarriers);
		
		//this is where you should screen out synergy only		
		if ( searcher.userParams.getSynergyOnly() ) {
			double[] disequilibria = getCaseAndControlDisequilibria(expectedCaseCarriers, noOfCaseCarriers, expectedControlCarriers, noOfControlCarriers);
			if ( disequilibria[0] - disequilibria[1] < 0 ) 
				return;
		}			
		
		try
		{
			switch( searcher.userParams.getVerboseState() ) 
			{
				case 0:
					WriteOutputBrief ( variableArray, 
							expectedCaseCarriers, noOfCaseCarriers, pvalueCaseOnly, 
							expectedControlCarriers, noOfControlCarriers, pvalueCaseControl ) ;
					break;
					
				case 1:
					WriteOutputNormal ( variableArray, 
							expectedCaseCarriers, noOfCaseCarriers, pvalueCaseOnly, 
							expectedControlCarriers, noOfControlCarriers, pvalueCaseControl ) ;
					break;
					
				default:
					WriteOutputFull ( variableArray, 
							expectedCaseCarriers, noOfCaseCarriers, pvalueCaseOnly, 
							expectedControlCarriers, noOfControlCarriers, pvalueCaseControl ) ;
					break;
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
		
	}
	
	public void writeOutputMatrix (double[][] pvalueMatrix, String[][] rowAndColLabels) throws IOException {
	
		//first line - write the column labels 
		bwriter.write("\t");
		for(int col=0; col<rowAndColLabels[1].length; col++) {
			bwriter.write(rowAndColLabels[1][col]);
			if(col!=rowAndColLabels[1].length-1)
				bwriter.write("\t");
			else
				bwriter.write("\n");
		}
		
		for(int i=0; i<pvalueMatrix.length; i++) {
			bwriter.write(rowAndColLabels[0][i]+"\t"); //row label
			for(int j=0; j<pvalueMatrix[i].length; j++) {
				bwriter.write( -Math.log10( pvalueMatrix[i][j] )+ "\t");
			}
			bwriter.write("\n");	
		}
		
	}
	
	
	public synchronized void WriteOutputClustered ( int[] variableArray, 
																						int[][][] noOfCaseCarriers, double pvalueCaseOnly, 
																						int[][][] noOfControlCarriers, double pvalueCaseControl,
																						Integer countOfNearbyHits ) 
	throws IOException { 
		
		double[][][] expectedCaseCarriers = computeExpectations(noOfCaseCarriers);
		double[][][] expectedControlCarriers = computeExpectations(noOfControlCarriers);
		
		int order = variableArray.length;
		String interactionString = getSNPDetailsfromVariables(variableArray);
		double[] disequilibria = getCaseAndControlDisequilibria(expectedCaseCarriers, noOfCaseCarriers, expectedControlCarriers, noOfControlCarriers);

		switch(order) {
		case 2:
			bwriter.write( interactionString.toString() 
					+ doubleF.format(expectedCaseCarriers[1][0][0]) + "/" + doubleF.format(expectedCaseCarriers[1][0][1]) + "/" + doubleF.format(expectedCaseCarriers[1][1][0]) + "/" + doubleF.format(expectedCaseCarriers[1][1][1]) + "\t" 
					+ noOfCaseCarriers[1][0][0]+"/"+noOfCaseCarriers[1][0][1]+"/"+noOfCaseCarriers[1][1][0]+"/"+noOfCaseCarriers[1][1][1] +"\t" 
					+ scientificF.format(pvalueCaseOnly) + "\t" 
					+ doubleF.format(expectedControlCarriers[1][0][0]) + "/" + doubleF.format(expectedControlCarriers[1][0][1]) + "/" + doubleF.format(expectedControlCarriers[1][1][0]) + "/" + doubleF.format(expectedControlCarriers[1][1][1]) + "\t" 
					+ noOfControlCarriers[1][0][0]+"/"+noOfControlCarriers[1][0][1]+"/"+ noOfControlCarriers[1][1][0]+"/"+noOfControlCarriers[1][1][1]+ "\t"
					+ scientificF.format(pvalueCaseControl) + "\t"
					+ twiceDoubleF.format(disequilibria[0]) + "\t"
					+ twiceDoubleF.format(disequilibria[1]) + "\t" 
					+ countOfNearbyHits + "\n" ); 
			break;
			
		case 3:
			bwriter.write ( interactionString.toString() +
					expectedCaseCarriers[0][0][0]+"/"+expectedCaseCarriers[0][0][1]+"/"+expectedCaseCarriers[0][1][0]+"/"+expectedCaseCarriers[0][1][1]+expectedCaseCarriers[1][0][0]+"/"+expectedCaseCarriers[1][0][1]+"/"+expectedCaseCarriers[1][1][0]+"/"+expectedCaseCarriers[1][1][1] + "\t" 
					+ noOfCaseCarriers[0][0][0]+"/"+noOfCaseCarriers[0][0][1]+"/"+noOfCaseCarriers[0][1][0]+"/"+noOfCaseCarriers[0][1][1]+ noOfCaseCarriers[1][0][0]+"/"+noOfCaseCarriers[1][0][1]+"/"+noOfCaseCarriers[1][1][0]+"/"+noOfCaseCarriers[1][1][1]  + "\t" 
					+ pvalueCaseOnly + "\t" 
					+ expectedControlCarriers[0][0][0]+"/"+expectedControlCarriers[0][0][1]+"/"+expectedControlCarriers[0][1][0]+"/"+expectedControlCarriers[0][1][1] +expectedControlCarriers[1][0][0]+"/"+expectedControlCarriers[1][0][1]+"/"+expectedControlCarriers[1][1][0]+"/"+expectedControlCarriers[1][1][1]  + "\t" 
					+ noOfControlCarriers[0][0][0]+"/"+noOfControlCarriers[0][0][1]+"/"+ noOfControlCarriers[0][1][0]+"/"+noOfControlCarriers[0][1][1] + noOfControlCarriers[1][0][0]+"/"+noOfControlCarriers[1][0][1]+"/"+ noOfControlCarriers[1][1][0]+"/"+noOfControlCarriers[1][1][1] + "\t"
					+ pvalueCaseControl + "\t"
					+ twiceDoubleF.format(disequilibria[0]) + "\t"
					+ twiceDoubleF.format(disequilibria[1]) + "\t" 
					+ countOfNearbyHits + "\n" ); 
			break;
			
		}
		bwriter.flush();
	}
	
	
	public synchronized void Close() throws IOException {
		bwriter.flush();
		bwriter.close();
	}
	
	
	
}
